package com.getan.boot.web.trace;

import com.getan.boot.common.ResponseData;
import com.getan.boot.common.holder.trace.TraceContextHolder;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.MethodParameter;
import org.springframework.web.context.request.NativeWebRequest;
import org.springframework.web.method.support.HandlerMethodReturnValueHandler;
import org.springframework.web.method.support.ModelAndViewContainer;
import org.springframework.web.servlet.mvc.method.annotation.RequestResponseBodyMethodProcessor;

/**
 * 渲染 TraceId
 *
 * @author Bruce.Gong
 * @since 1.0.0-SNAPSHOT
 */
@Slf4j
@RequiredArgsConstructor
public class TraceResponseHandler implements HandlerMethodReturnValueHandler {

    private final RequestResponseBodyMethodProcessor target;

    @Override
    public boolean supportsReturnType(MethodParameter returnType) {
        return target.supportsReturnType(returnType);
    }

    @Override
    public void handleReturnValue(Object returnValue, MethodParameter returnType, ModelAndViewContainer mavContainer, NativeWebRequest webRequest) throws Exception {
        if (returnValue instanceof ResponseData) {
            ((ResponseData<?>) returnValue).setTraceId(TraceContextHolder.getTraceId());
        }
        target.handleReturnValue(returnValue, returnType, mavContainer, webRequest);
    }

}
